package main

import (
	"fmt"
	"sort"
	"strconv"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/pluginpb"
)

// generate generates Go code with comment registries from protobuf descriptors
func generate(req *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
	opts := protogen.Options{}
	plugin, err := opts.New(req)
	if err != nil {
		return nil, fmt.Errorf("failed to create plugin: %w", err)
	}

	plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)

	// Group files by package to generate one comments file per package
	packageFiles := make(map[protogen.GoImportPath][]*protogen.File)
	for _, file := range plugin.Files {
		if !file.Generate {
			continue
		}
		if len(file.Messages) == 0 && len(file.Enums) == 0 {
			continue
		}
		packageFiles[file.GoImportPath] = append(packageFiles[file.GoImportPath], file)
	}

	// Generate one comments file per package
	for importPath, files := range packageFiles {
		generatePackageCommentsFile(plugin, files, importPath)
	}

	return plugin.Response(), nil
}

// generatePackageCommentsFile generates a single comments.pb.go file for all proto files in a package
func generatePackageCommentsFile(plugin *protogen.Plugin, files []*protogen.File, importPath protogen.GoImportPath) {
	if len(files) == 0 {
		return
	}

	// Use the first file to determine package name
	firstFile := files[0]

	// With paths=import, derive filename from import path
	// The import path is the full Go import path, and we need to extract
	// the relative path from the module root for the filename
	importPathStr := string(importPath)

	// Extract everything after "github.com/redpanda-data/redpanda/"
	const modulePrefix = "github.com/redpanda-data/redpanda/"
	var filename string
	if idx := strings.Index(importPathStr, modulePrefix); idx >= 0 {
		// Get the relative path, e.g: src/go/rpk/gen/protocomments/admin/v2
		relativePath := importPathStr[idx+len(modulePrefix):]
		filename = relativePath + "/comments.pb.go"
	} else {
		// Fallback: use a simple filename (may cause conflicts)
		filename = "comments.pb.go"
	}

	g := plugin.NewGeneratedFile(filename, importPath)

	// Generate file header
	g.P("// Code generated by protoc-gen-comments. DO NOT EDIT.")
	g.P("// This file consolidates comments from all proto files in this package.")
	g.P()
	g.P("package ", firstFile.GoPackageName)
	g.P()

	// Collect all comments from all files in this package
	allComments := make(map[string]string)
	for _, file := range files {
		fileComments := buildCommentMapByFullName(file)
		for fullName, comment := range fileComments {
			allComments[fullName] = comment
		}
	}

	// Generate package-level comment map (sorted for deterministic output)
	keys := make([]string, 0, len(allComments))
	for k := range allComments {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// Write comments map using g.P() but with raw strings for multi-line comments
	g.P("// Comments maps protobuf descriptor full names to their leading comments")
	g.P("var Comments = map[string]string{")
	for _, fullName := range keys {
		comment := allComments[fullName]
		// If comment contains newlines, use raw string literal (backticks)
		// Otherwise use regular quoted string
		var quotedComment string
		hasNewline := strings.Contains(comment, "\n")
		if hasNewline {
			// Use backtick raw string for multi-line comments
			quotedComment = "`" + strings.ReplaceAll(comment, "`", "` + \"`\" + `") + "`"
		} else {
			quotedComment = strconv.Quote(comment)
		}
		// Build the line and write it with raw string template
		// to prevent g.P() from escaping
		line := fmt.Sprintf("\t%s: %s,", strconv.Quote(fullName), quotedComment)
		g.P(line)
	}
	g.P("}")
	g.P()

	// Generate simple helper function
	g.P("// GetComment returns the comment for a descriptor by its full name")
	g.P("func GetComment(fullName string) string {")
	g.P("\treturn Comments[fullName]")
	g.P("}")
	g.P()

	// Generate convenience helpers that take descriptors directly
	g.P("// GetCommentForField returns the comment for a field descriptor")
	g.P("func GetCommentForField(fd ", g.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "google.golang.org/protobuf/reflect/protoreflect", GoName: "FieldDescriptor"}), ") string {")
	g.P("\treturn GetComment(string(fd.FullName()))")
	g.P("}")
	g.P()

	g.P("// GetCommentForMessage returns the comment for a message descriptor")
	g.P("func GetCommentForMessage(md ", g.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "google.golang.org/protobuf/reflect/protoreflect", GoName: "MessageDescriptor"}), ") string {")
	g.P("\treturn GetComment(string(md.FullName()))")
	g.P("}")
	g.P()

	g.P("// GetCommentForEnum returns the comment for an enum descriptor")
	g.P("func GetCommentForEnum(ed ", g.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "google.golang.org/protobuf/reflect/protoreflect", GoName: "EnumDescriptor"}), ") string {")
	g.P("\treturn GetComment(string(ed.FullName()))")
	g.P("}")
	g.P()

	g.P("// GetCommentForEnumValue returns the comment for an enum value descriptor")
	g.P("func GetCommentForEnumValue(evd ", g.QualifiedGoIdent(protogen.GoIdent{GoImportPath: "google.golang.org/protobuf/reflect/protoreflect", GoName: "EnumValueDescriptor"}), ") string {")
	g.P("\treturn GetComment(string(evd.FullName()))")
	g.P("}")
}

// buildCommentMapByFullName builds a map from descriptor FullName to comments
func buildCommentMapByFullName(file *protogen.File) map[string]string {
	result := make(map[string]string)

	// Get source code info from the proto descriptor
	sourceCodeInfo := file.Proto.GetSourceCodeInfo()
	if sourceCodeInfo == nil {
		return result
	}

	// First, build a map from proto path to comments
	pathToComment := make(map[string]string)
	for _, loc := range sourceCodeInfo.Location {
		if loc.LeadingComments == nil || *loc.LeadingComments == "" {
			continue
		}
		path := pathToString(loc.Path)
		pathToComment[path] = strings.TrimSpace(*loc.LeadingComments)
	}

	// Now map descriptors to their FullNames
	// Process messages
	for _, msg := range file.Messages {
		processMessage(msg, pathToComment, result, []int32{4})
	}

	// Process enums
	for _, enum := range file.Enums {
		processEnum(enum, pathToComment, result, []int32{5})
	}

	return result
}

// processMessage recursively processes a message and its nested types
func processMessage(msg *protogen.Message, pathToComment map[string]string, result map[string]string, pathPrefix []int32) {
	// Build path for this message
	msgPath := append(pathPrefix, int32(msg.Desc.Index()))
	msgPathStr := pathToString(msgPath)

	// Store message comment
	if comment, ok := pathToComment[msgPathStr]; ok {
		result[string(msg.Desc.FullName())] = comment
	}

	// Process fields
	for _, field := range msg.Fields {
		fieldPath := append(msgPath, 2, int32(field.Desc.Index()))
		fieldPathStr := pathToString(fieldPath)
		if comment, ok := pathToComment[fieldPathStr]; ok {
			result[string(field.Desc.FullName())] = comment
		}
	}

	// Process nested enums
	for _, enum := range msg.Enums {
		enumPath := append(msgPath, 4)
		processEnum(enum, pathToComment, result, enumPath)
	}

	// Process nested messages
	for _, nested := range msg.Messages {
		nestedPath := append(msgPath, 3)
		processMessage(nested, pathToComment, result, nestedPath)
	}
}

// processEnum processes an enum and its values
func processEnum(enum *protogen.Enum, pathToComment map[string]string, result map[string]string, pathPrefix []int32) {
	// Build path for this enum
	enumPath := append(pathPrefix, int32(enum.Desc.Index()))
	enumPathStr := pathToString(enumPath)

	// Store enum comment
	if comment, ok := pathToComment[enumPathStr]; ok {
		result[string(enum.Desc.FullName())] = comment
	}

	// Process enum values
	for _, value := range enum.Values {
		valuePath := append(enumPath, 2, int32(value.Desc.Index()))
		valuePathStr := pathToString(valuePath)
		if comment, ok := pathToComment[valuePathStr]; ok {
			result[string(value.Desc.FullName())] = comment
		}
	}
}

// pathToString converts a protobuf path to a string representation
func pathToString(path []int32) string {
	if len(path) == 0 {
		return ""
	}
	parts := make([]string, len(path))
	for i, p := range path {
		parts[i] = fmt.Sprintf("%d", p)
	}
	return strings.Join(parts, ".")
}
